from biz.data.sql_template.sql_generator import SQLGenerator
from biz.data.sql_template.sql_template_store import SqlTemplateStore


class SqlTemplateRAG:
    def __init__(self,config=None, openLLM=None):
        if openLLM is None:
            raise ValueError("openLLM is required")

        self.template_store = SqlTemplateStore(config)
        self.sql_generator = SQLGenerator(openLLM)

    def create_function(self, id, description, sql_template, parameters, plotly_code=None):
        """创建新的SQL函数模板"""
        self.template_store.add_sql_template(
            id=id,
            description=description,
            sql_template=sql_template,
            parameters=parameters,
            plotly_code=plotly_code
        )
        print(f"Function template '{id}' created successfully.")
        return id

    def update_function(self, id, description=None, sql_template=None, parameters=None, plotly_code=None):
        """更新SQL函数模板"""
        success = self.template_store.update_sql_template(
            id=id,
            description=description,
            sql_template=sql_template,
            parameters=parameters,
            plotly_code=plotly_code
        )
        if success:
            print(f"Function template '{id}' updated successfully.")
            return True
        else:
            print(f"Failed to update function template '{id}'.")
            return False

    def delete_function(self, id):
        """删除SQL函数模板"""
        return self.template_store.delete_sql_template(id)

    def get_function_by_id(self, id):
        """根据ID获取SQL函数模板"""
        template = self.template_store.get_sql_template_by_id(id)
        if not template:
            print(f"No function template found with ID: {id}")
            return None
        return template

    def search_similar_templates(self, query):
        """获取相似性模板"""
        templates = self.template_store.search_similar_templates(query)
        return templates

    def generate_sql_from_question(self, question, additional_data=None):
        """
        根据用户问题搜索模板并生成SQL

        Args:
            question: 用户问题
            additional_data: 额外数据，如user_id等

        Returns:
            包含生成SQL和可选Plotly代码的字典
        """
        # 获取相似的模板
        templates = self.template_store.search_similar_templates(question)

        if not templates:
            return {"error": "No templates available"}

        # 生成SQL
        try:
            sql_json = self.sql_generator.generate_sql(templates, question)
            plotly_code = self.sql_generator.generate_plotly_code(templates,sql_json['matched_question_id'],)

            result = {
                "sql": sql_json['sql'],
            }

            if plotly_code:
                result["plotly_code"] = plotly_code

            return result
        except Exception as e:
            return {"error": str(e)}

    def list_functions(self):
        """列出所有SQL函数模板"""
        return self.template_store.list_all_templates()